# -*- coding: utf-8 -*-

import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler, OneHotEncoder
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import TensorDataset, DataLoader
import matplotlib.pyplot as plt
from sklearn.metrics import precision_score, recall_score, confusion_matrix, accuracy_score, f1_score
from torchviz import make_dot

plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']  # 设置字体为微软雅黑

# 设置随机种子以确保结果可复现
np.random.seed(42)
torch.manual_seed(42)

from gen_corr import get_corr_feature

# 加载数据
train_df = pd.read_csv('data_c1.csv')
test_df_4 = pd.read_csv('data_c4.csv')
test_df_6 = pd.read_csv('data_c6.csv')
cols = get_corr_feature()

# 分离特征和标签
X_train = train_df.drop('label', axis=1)[cols].values
y_train = train_df['label'].values

# 分离特征和标签
X_test_4 = test_df_4.drop('label', axis=1)[cols].values
y_test_4 = test_df_4['label'].values

# 分离特征和标签
X_test_6 = test_df_6.drop('label', axis=1)[cols].values
y_test_6 = test_df_6['label'].values

# 数据归一化
scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_test_4 = scaler.fit_transform(X_test_4)
X_test_6 = scaler.fit_transform(X_test_6)

# 将标签转换为one-hot编码
encoder = OneHotEncoder(sparse=False)
y_train = encoder.fit_transform(y_train.reshape(-1, 1))
y_test_4 = encoder.fit_transform(y_test_4.reshape(-1, 1))
y_test_6 = encoder.fit_transform(y_test_6.reshape(-1, 1))

# 转换为PyTorch tensors
X_train = torch.tensor(X_train, dtype=torch.float).unsqueeze(1)
y_train = torch.tensor(y_train, dtype=torch.float)
X_test_4 = torch.tensor(X_test_4, dtype=torch.float).unsqueeze(1)
y_test_4 = torch.tensor(y_test_4, dtype=torch.float)
X_test_6 = torch.tensor(X_test_6, dtype=torch.float).unsqueeze(1)
y_test_6 = torch.tensor(y_test_6, dtype=torch.float)
# 创建数据加载器
batch_size = 32
train_loader = DataLoader(TensorDataset(X_train, y_train), batch_size=batch_size, shuffle=True)


# test_loader_4 = DataLoader(TensorDataset(X_test_4, y_test_4), batch_size=batch_size)
# test_loader_6 = DataLoader(TensorDataset(X_test_6, y_test_4), batch_size=batch_size)


# 定义1D CNN模型
class CNN1d(nn.Module):
    def __init__(self):
        super(CNN1d, self).__init__()
        self.conv1 = nn.Conv1d(1, 32, 3, padding=1)
        self.conv2 = nn.Conv1d(32, 64, 3, padding=1)
        self.conv3 = nn.Conv1d(64, 128, 3, padding=1)
        self.pool = nn.MaxPool1d(2, 2)
        self.fc1 = nn.Linear(128 * (X_train.shape[2] // 8), 128)  # 注意适当调整以匹配维度
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, 3)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = self.pool(self.relu(self.conv2(x)))
        x = self.pool(self.relu(self.conv3(x)))
        x = x.view(-1, 128 * (X_train.shape[2] // 8))
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        x = self.fc3(x)
        return x


# 实例化模型、损失函数和优化器
model = CNN1d()
# 加载模型权重（选择你需要的路径）
model.load_state_dict(torch.load('best_model_4.pth'))
# 将模型置于评估模式
model.eval()
# 无需计算梯度
with torch.no_grad():
    # 进行预测
    outputs = model(X_test_4)
    # 获取每个样本最可能的类别索引
    _, predicted_indices = torch.max(outputs, 1)
pre_4 = [t.item() for t in predicted_indices]
true_4 = list(test_df_4['label'].values)
print(pre_4)
print(true_4)
# 创建一个范围列表，确保曲线的X轴对应（如果你有特定的X轴数据，也可以使用那些）
x_range = range(len(pre_4))

# 绘制真实类别曲线
plt.plot(x_range, true_4, label='真实类别', marker='o')  # 'o'代表圆圈标记

# 绘制预测类别曲线
plt.plot(x_range, pre_4, label='预测类别', marker='x')  # 'x'代表叉号标记

# 添加图例
plt.legend()

# 添加标题和轴标签
plt.title('刀具4测试集1DCNN预测结果对比')
plt.xlabel('刀具4测试集样本编号')
plt.ylabel('刀具4测试集样本类别')
plt.savefig('cnn_1.png', dpi=300)

plt.clf()
model.load_state_dict(torch.load('best_model_6.pth'))
# 将模型置于评估模式
model.eval()
# 无需计算梯度
with torch.no_grad():
    # 进行预测
    outputs = model(X_test_6)
    # 获取每个样本最可能的类别索引
    _, predicted_indices = torch.max(outputs, 1)
pre_6 = [t.item() for t in predicted_indices]
true_6 = list(test_df_6['label'].values)
print(pre_6)
print(true_6)
# 创建一个范围列表，确保曲线的X轴对应（如果你有特定的X轴数据，也可以使用那些）
x_range = range(len(pre_6))

# 绘制真实类别曲线
plt.plot(x_range, true_6, label='真实类别', marker='o')  # 'o'代表圆圈标记

# 绘制预测类别曲线
plt.plot(x_range, pre_6, label='预测类别', marker='x')  # 'x'代表叉号标记

# 添加图例
plt.legend()

# 添加标题和轴标签
plt.title('刀具6测试集1DCNN预测结果对比')
plt.xlabel('刀具6测试集样本编号')
plt.ylabel('刀具6测试集样本类别')
plt.savefig('cnn_2.png', dpi=300)

import seaborn as sns

# 计算刀具4的混淆矩阵
cm_4 = confusion_matrix(true_4, pre_4)
# 计算刀具6的混淆矩阵
cm_6 = confusion_matrix(true_6, pre_6)
class_names = ['初期磨损', '初期磨损', '急剧磨损']
# 绘制刀具4的混淆矩阵
plt.figure(figsize=(10, 7))
sns.heatmap(cm_4, annot=True, fmt="d", cmap="Blues", xticklabels=class_names, yticklabels=class_names)
plt.title('1DCNN刀具4测试集混淆矩阵')
plt.xlabel('预测磨损状态')
plt.ylabel('真实磨损状态')
plt.savefig('cnn_confusion_matrix_4.png', dpi=300)
plt.clf()  # 清空画布以便绘制下一个混淆矩阵

# 绘制刀具6的混淆矩阵
plt.figure(figsize=(10, 7))
sns.heatmap(cm_6, annot=True, fmt="d", cmap="Blues", xticklabels=class_names, yticklabels=class_names)
plt.title('1DCNN刀具6测试集混淆矩阵')
plt.xlabel('预测磨损状态')
plt.ylabel('真实磨损状态')
plt.savefig('cnn_confusion_matrix_6.png', dpi=300)

# 计算准确率、精度和召回率
accuracy_4 = accuracy_score(true_4, pre_4)
precision_4 = precision_score(true_4, pre_4, average='macro')
recall_4 = recall_score(true_4, pre_4, average='macro')
f1_4 = f1_score(true_4, pre_4, average='macro')

# 计算准确率、精度和召回率
accuracy_6 = accuracy_score(true_6, pre_6)
precision_6 = precision_score(true_6, pre_6, average='macro')  # 确保这里的average参数符合你的分类问题
recall_6 = recall_score(true_6, pre_6, average='macro')
f1_6 = f1_score(true_6, pre_6, average='macro')
# 输出结果
# 将指标写入到文本文件
with open('classification_report.txt', 'a') as f:
    f.write('CNN')
    f.write(
        f'\nTest Set 4 - Accuracy: {accuracy_4:.4f}, Precision: {precision_4:.4f}, Recall: {recall_4:.4f}, F1 Score: {f1_4:.4f}')
    f.write(
        f'\nTest Set 6 - Accuracy: {accuracy_6:.4f}, Precision: {precision_6:.4f}, Recall: {recall_6:.4f}, F1 Score: {f1_6:.4f}\n')
